package com.ymf.aiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.ymf.aiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * 抽象基础代理类，用于管理代理状态和执行流程。
 * <p>
 * 提供状态转换、内存管理和基于步骤的执行循环的基础功能。
 * 子类必须实现step方法。
 */
@Data
@Slf4j
public abstract class BaseAgent {
    // 核心属性
    private String name;

    // 提示词
    private String systemPrompt;
    private String nextStepPrompt;

    // 代理状态
    private AgentState state = AgentState.IDLE;

    //执行步骤控制
    private int currentStep = 0;
    private int maxSteps = 10;//最大步骤

    // LLM大模型
    private ChatClient chatClient;


    // Memory 记忆（需要自主维护会话上下文）
    private List<Message> messageList = new ArrayList<>();

    /**
     * 运行代理
     *
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public String run(String userPrompt) {

        // 1.先校验运行状态
        if (this.state != AgentState.IDLE) {
            throw new RuntimeException("Cannot run agent from state: " + this.state);
        }
        //用户提示词不能为空
        if (!StringUtils.hasText(userPrompt)) {
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }
        // 2.更改执行状态
        this.state = AgentState.RUNNING;
        //记录消息上下文
        messageList.add(new UserMessage(userPrompt));
        //保存结果的列表
        ArrayList<String> resultList = new ArrayList<>();

        // 3.执行循环

        try {
            for (int i = 0; i < maxSteps && this.state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Executing step {}/{}", stepNumber, maxSteps);
                // 单步执行
                String stepResult = this.step();
                String result = "Step " + stepNumber + ": " + stepResult;
                resultList.add(result);
            }
            // 检查是否超出步骤限制
            if (currentStep >= maxSteps) {
                state = AgentState.FINISHED;
                resultList.add("Terminated: Reached max steps (" + maxSteps + ")");
            }
            return String.join("\n", resultList);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("error executing agent", e);
            return "执行错误" + e.getMessage();
        } finally {
            //清理资源
            this.cleanup();
        }

    }

    /**
     * 运行代理（流式输出）
     *
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public SseEmitter runStream(String userPrompt) {

        SseEmitter sseEmitter = new SseEmitter(300000L);//超时时长5分钟


        CompletableFuture.runAsync(()->{
            try {
                // 1.先校验运行状态
                if (this.state != AgentState.IDLE) {
                    sseEmitter.send("错误：无法从状态运行代理：" + this.state);
                    sseEmitter.complete();
                    return;
                }
                //用户提示词不能为空
                if (StrUtil.isBlank(userPrompt)) {
                    sseEmitter.send("错误：不能使用空提示词运行代理");
                    sseEmitter.complete();
                    return;
                }
                // 2.更改执行状态
                this.state = AgentState.RUNNING;
                //记录消息上下文
                messageList.add(new UserMessage(userPrompt));
                //保存结果的列表
                ArrayList<String> resultList = new ArrayList<>();

                // 3.执行循环

                try {
                    for (int i = 0; i < maxSteps && this.state != AgentState.FINISHED; i++) {
                        int stepNumber = i + 1;
                        currentStep = stepNumber;
                        log.info("Executing step {}/{}", stepNumber, maxSteps);
                        // 单步执行
                        String stepResult = this.step();
                        String result = "Step " + stepNumber + ": " + stepResult;
                        resultList.add(result);
                        // 输出当前每一步的结果到 SSE
                        sseEmitter.send(result);
                    }
                    // 检查是否超出步骤限制
                    if (currentStep >= maxSteps) {
                        state = AgentState.FINISHED;
                        resultList.add("Terminated: Reached max steps (" + maxSteps + ")");
                        // 输出当前每一步的结果到 SSE
                        sseEmitter.send("执行结束：达到最大步骤（" + maxSteps + "）");
                    }
                    // 正常完成
                    sseEmitter.complete();
                } catch (Exception e) {
                    state = AgentState.ERROR;
                    log.error("error executing agent", e);
                    try {
                        sseEmitter.send("执行错误：" + e.getMessage());
                        sseEmitter.complete();
                    } catch (IOException ex) {
                        sseEmitter.completeWithError(ex);
                    }
                } finally {
                    //清理资源
                    this.cleanup();
                }
            } catch (Exception e) {
                log.error("error executing agent", e);
                sseEmitter.completeWithError(e);
            }
        });

        return sseEmitter;


    }

    /**
     * 定义单个步骤
     *
     * @return
     */
    public abstract String step();

    /**
     * 清理资源
     */
    protected void cleanup() {
        // 子类可以重写此方法来清理资源
    }

}
